import os
import sys
import pandas as pd
from . import file_util


def output_excel(output_dir, dataframes_dict, file_name, merge_sheet=False, merge_sheet_name="Sheet1"):
    """
    output_dir: 文件路径
    df: DataFrame
    file_name: 文件名
    merge_sheet: 是否合并所有 sheet 到一个 sheet 中
    merge_sheet_name: 合并后的 sheet 名称
    """

    # 确保输出目录存在
    file_util.safe_mkdir(output_dir)
    out_path = os.path.join(output_dir, file_name)

    if merge_sheet:
        dataframes_dict.to_excel(out_path, sheet_name=merge_sheet_name, index=False)
    else:
        # 使用 ExcelWriter 来写入多个 sheet
        with pd.ExcelWriter(out_path) as writer:
            for sheet_name, data_frame in dataframes_dict.items():
                data_frame.to_excel(writer, sheet_name=sheet_name, index=False)

    print(f"已保存 Excel 到：{out_path}")


def input_excel(file_path, sheet_name):
    """
    file_path: 文件路径
    sheet_name: 表名
    """
    try:
        return pd.read_excel(file_path, sheet_name=sheet_name)
    except Exception as e:
        print(f"无法读取 Excel 文件：{file_path}（请确认路径和 sheet 名称）。错误：{e}")
        sys.exit(1)
